import tensorflow as tf
import random
from pprint import pprint as P
from keras.models import load_model
from pathlib import Path
import os
import matplotlib.pyplot as plt
import pandas as pd
from statsmodels.tsa.stattools import pacf
import numpy as np
from sklearn import preprocessing
from keras.layers.recurrent import LSTM, GRU
from keras.models import Sequential
from keras.layers import Dense, Activation


def model_train(mode_type, epochs, x_train, y_train, x_test, y_test, x_dim, y_dim):
    seed = 2022
    np.random.seed(seed)
    if mode_type == 'GRU':
        model = Sequential()
        model.add(GRU(units=4,
                      return_sequences=False,
                      activation='tanh',
                      recurrent_activation='hard_sigmoid',
                      input_shape=(5, 1)))
        model.add(Dense(units=1, activation='linear'))
        model.compile(loss='mean_squared_error', optimizer='rmsprop')
    else:
        model = Sequential()
        model.add(LSTM(units=4, activation='tanh', recurrent_activation='hard_sigmoid', input_shape=(5, 1)))
        model.add(Dense(units=1, activation='linear'))
        model.compile(loss='mean_squared_error', optimizer='rmsprop')

    model.fit(x_train, y_train, batch_size=1, epochs=epochs, shuffle=False)
    print(model.summary())

    score_train = model.evaluate(x_train, y_train, batch_size=1)
    score_test = model.evaluate(x_test, y_test, batch_size=1)
    print("in train MSE = ", round(score_train, 4))
    print("in test MSE = ", round(score_test, 4))
    pred = model.predict(x_test)
    plt.plot(y_test[:, 0])
    plt.plot(pred[:, 0])
    plt.legend([f'{mode_type}_t', f'{mode_type}_p1'])
    plt.show()
    return model


def load_data(file_path, col_index):
    # 加载数据集合
    data = pd.read_csv(file_path)

    # 拉取某一列数据
    y = data.iloc[:, col_index]
    # x_pacf = pacf(y, nlags=20, method='ols')
    # plt.plot(x_pacf)
    # plt.show()
    return y


def build_rect(row, x_dim=5, y_dim=1):
    yt = [row, ]
    yt_shifts = [row.shift(i) for i in range(1, x_dim+y_dim)]
    yt.extend(yt_shifts)
    data = pd.concat(yt, axis=1)
    data = data.iloc[x_dim+y_dim:]
    data.columns = [f'yt_{i}' for i in range(1, x_dim+y_dim+1)]
    data = data.dropna()
    y = data[[f'yt_{i}' for i in range(1, y_dim+1)]]
    x = data[[f'yt_{i}' for i in range(y_dim, y_dim+x_dim)]]
    return x, y


def normalize(x, dim):
    scaler = preprocessing.MinMaxScaler(feature_range=(-1, 1))
    x_1d = x.to_numpy().reshape(len(x), dim)
    x_norm = scaler.fit_transform(x_1d)
    return x_norm, scaler


def normal(x, y):
    # 归一化
    scaler_x = preprocessing.MinMaxScaler(feature_range=(-1, 1))
    x = np.array(x).reshape((len(x), 5))
    x = scaler_x.fit_transform(x)
    scaler_y = preprocessing.MinMaxScaler(feature_range=(-1, 1))
    y = np.array(y).reshape((len(y), 1))
    y = scaler_y.fit_transform(y)
    train_end = int(len(x) * 0.95)
    x_train = x[0: train_end, ]
    x_test = x[train_end:]
    y_train = y[0: train_end]
    y_test = y[train_end:]
    x_train = x_train.reshape(x_train.shape + (1,))
    x_test = x_test.reshape(x_test.shape + (1,))
    print(x_train.shape)
    return x_train, y_train, x_test, y_test, scaler_x, scaler_y


def main(model_type):
    data = load_data('../data/XAUUSD_txt.csv.day', 4)

    x_dim = 64
    y_dim = 10
    x_rect, y_rect = build_rect(data, x_dim, y_dim)

    x_norm_rect, x_scaler = normalize(x_rect, x_dim)
    y_norm_rect, y_scaler = normalize(y_rect, y_dim)

    train_end = int(len(x_rect) * 0.95)
    x_train = x_norm_rect[0: train_end, ]
    x_test = x_norm_rect[train_end:]

    y_train = y_norm_rect[0: train_end]
    y_test = y_norm_rect[train_end:]
    model = model_train(
        model_type, 1,
        x_train.reshape(x_train.shape + (1,)),
        y_train.reshape(y_train.shape + (1,)),
        x_test.reshape(x_test.shape + (1,)),
        y_test.reshape(y_test.shape + (1,)),
        x_dim, y_dim
    )

    y_pred = model.predict(x_test.reshape(x_test.shape + (1,)))
    result = y_scaler.inverse_transform(y_pred)
    plt.plot(result[:, 0])
    plt.legend([f'pred'])
    plt.show()


if __name__ == "__main__":
    main("GRU")
